package pl.edu.agh.broker.mns.server.internal;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;

public class MultiMap<K, V> implements Iterable<V> {

	private Map<K, Collection<V>> map;

	public MultiMap() {
		this.map = new HashMap<K, Collection<V>>();
	}

	public Collection<V> get(K key) {
		Collection<V> current = map.get(key);
		if (current == null) {
			return Collections.emptySet();
		}
		return current;
	}

	public void put(K key, V value) {
		Collection<V> current = map.get(key);
		if (current == null) {
			current = new HashSet<V>();
			map.put(key, current);
		}
		current.add(value);
	}

	public void remove(K key) {
		map.remove(key);
	}

	public void remove(K key, V value) {
		Collection<V> c = map.get(key);
		if (c != null) {
			c.remove(value);
			if (c.isEmpty()) {
				map.remove(key);
			}
		}
	}

	@Override
	public Iterator<V> iterator() {
		return new Iterator<V>() {

			private Iterator<Collection<V>> collections = map.values()
					.iterator();
			private Iterator<V> elements;
			private V next;

			@Override
			public boolean hasNext() {
				next = null;
				if (elements == null) {
					if (collections.hasNext()) {
						elements = collections.next().iterator();
					}
				}
				if (elements != null && elements.hasNext()) {
					next = elements.next();
				}
				return next != null;
			}

			@Override
			public V next() {
				return next;
			}

			@Override
			public void remove() {
				throw new UnsupportedOperationException();
			}
		};
	}
}
